package online.heycm.platform.web.common.sse;

import java.io.IOException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import lombok.extern.slf4j.Slf4j;
import online.heycm.platform.common.tools.Jackson;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

/**
 * SSE 客户端
 *
 * @author heycm
 * @since 2023/12/16 0:08
 */
@Slf4j
public class SseSession {

    /**
     * SSE ID
     */
    private final String id;

    /**
     * 推送器
     */
    private final SseEmitter emitter;

    /**
     * 当前推送器是否存活
     */
    private final AtomicBoolean isAlive = new AtomicBoolean(true);

    /**
     * 推送消息计数
     */
    private final AtomicInteger row = new AtomicInteger(0);

    public SseSession(String id) {
        this.id = id;
        this.emitter = new SseEmitter(0L);
    }

    public SseSession(String id, long timeout) {
        this.id = id;
        this.emitter = new SseEmitter(timeout);
    }

    public String getId() {
        return id;
    }

    public SseEmitter getSseEmitter() {
        return emitter;
    }

    /**
     * 当前连接是否存活
     */
    public boolean isAlive() {
        return isAlive.get();
    }

    /**
     * 关闭连接
     */
    public void close() {
        log.info("SSE[{}] 关闭连接", id);
        if (isAlive()) {
            isAlive.set(false);
            emitter.complete();
        }
    }

    /**
     * 发送消息
     *
     * @param content 消息内容
     */
    public boolean send(String content) {
        return send(SseSet.MESSAGE, content);
    }

    /**
     * 发送消息，指定接收事件
     *
     * @param event   事件名称
     * @param content 消息内容
     */
    public boolean send(String event, String content) {
        SseEvent sseEvent = new SseEvent();
        sseEvent.setEvent(event);
        sseEvent.setData(content);
        sseEvent.setRetry(SseSet.RETRY);
        return send(sseEvent);
    }

    /**
     * 发送事件
     *
     * @param event 事件
     */
    private boolean send(SseEvent event) {
        if (!isAlive()) {
            log.error("SSE[{}] 已失活，推送消息失败", id);
            return false;
        }
        if (event.getData() == null) {
            log.error("SSE[{}] 推送消息为null", id);
            return true;
        }
        if (!StringUtils.hasText(event.getId())) {
            event.setId(row.getAndIncrement() + "");
        }
        log.info("SSE[{}] 推送消息: {}", id, Jackson.toJson(event));
        try {
            emitter.send(event.sseEventBuilder());
            return true;
        } catch (IOException e) {
            log.info("SSE[{}] 推送消息发生异常", id, e);
            emitter.completeWithError(e);
            throw new RuntimeException(e);
        }
    }
}
